import numpy as np
import cv2
import json
from pathlib import Path
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import os
import argparse

from os.path import join, split, exists
import pdb 
def w2c_to_c2w(R_w2c, t_w2c):
  R_c2w = R_w2c.T  # 旋转矩阵的逆是其转置
  t_c2w = -np.dot(R_c2w, t_w2c)  # 平移向量的逆是旋转矩阵的转置乘以原始平移向量的负值
  return R_c2w, t_c2w
def project_to_image(pts_3d, K):
    pts_2d = K.dot(pts_3d.T).T
    pts_2d /= pts_2d[:, 2].reshape(-1, 1)
    return pts_2d[:, :2]
def draw_bbox(bbox, color, linewidth=2):
  x, y, w, h = bbox
  rect = patches.Rectangle((x, y), w, h, linewidth=linewidth, edgecolor=color, facecolor='none')
  plt.gca().add_patch(rect)

def data_prepare(root_path, output_path):
    camera_info_path = Path(root_path) / "scene_camera.json"
    image_info_path = Path(root_path) / "scene_gt_info.json"
    obj_info_path = Path(root_path) / "scene_gt.json"
    rgb_path = Path(root_path) / "rgb"
    with open(camera_info_path, "r") as f:
        camera_info = json.load(f)
    with open(image_info_path, "r") as f:
        image_info = json.load(f)
    with open(obj_info_path, "r") as f:
        obj_info = json.load(f)
    os.makedirs(output_path, exist_ok=True)
    return camera_info, image_info, obj_info, rgb_path

def pose_visualization(image, R, T, K, bbox_obj, bbox_visib, output_file, axis_length = 50):
    plt.imshow(image)
    axis = np.float32([[0, 0, 0],
                [axis_length, 0, 0],
                [0, axis_length, 0],
                [0, 0, axis_length]])
    axis_cam = np.dot(R, axis.T).T + T
    imgpts = project_to_image(axis_cam, K)
    colors = [(1, 0, 0),
            (0, 1, 0),
            (0, 0, 1)]
    origin = imgpts[0]
    for i, col in enumerate(colors):
        end_point = imgpts[i + 1]
        plt.plot([origin[0], end_point[0]], [origin[1], end_point[1]], color=col, linewidth=2, label=f'Axis {i}')
    draw_bbox(bbox_obj, "yellow")
    draw_bbox(bbox_visib, "black")
    plt.axis('off')
    plt.savefig(output_file)    
    plt.clf()

def main(root_path, output_path):
    camera_info, image_info, obj_info, rgb_path = data_prepare(root_path, output_path)
    num_images = len(list(rgb_path.glob("*.jpg")))
    plt.figure(figsize=(6.4, 4.8))
    for idx in range(num_images):
        file_name = f"{idx:0>6d}.jpg"
        rgb = Path(rgb_path) / file_name
        
        image = cv2.imread(str(rgb))[:,:,::-1]
        idx = str(idx)
        K = np.array(camera_info[idx]["cam_K"]).reshape([3,3])
        R = np.array(obj_info[idx][0]["cam_R_m2c"]).reshape([3,3])
        T = np.array(obj_info[idx][0]["cam_t_m2c"])
        bbox_obj = np.array(image_info[idx][0]["bbox_obj"])
        bbox_visib = np.array(image_info[idx][0]["bbox_visib"])
        output_file = output_path / file_name
        pose_visualization(image, R, T, K, bbox_obj, bbox_visib, output_file)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--root_path', type=Path, default="/baai-cwm-1/baai_cwm_ml/cwm/shaocong.xu/exp/BlenderProc/examples/datasets/bop_challenge/output_debug_front5/bop_data/lm/train_pbr/000000")
    parser.add_argument('--output_path', type=Path, default="logs/pose")
    args = parser.parse_args()

    root_path = args.root_path
    output_path = args.output_path
    if not exists(output_path):
        os.makedirs(output_path)

    main(root_path, output_path)